package com.zenithmind.news.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.time.LocalDateTime;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * 新闻推送WebSocket处理器
 * 
 * @author ZenithMind Team
 * @since 2025-01-09
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class NewsWebSocketHandler implements WebSocketHandler {

    private final ObjectMapper objectMapper;

    // 存储所有WebSocket连接 - 用户ID -> WebSocket会话集合
    private static final Map<String, CopyOnWriteArraySet<WebSocketSession>> USER_SESSIONS = new ConcurrentHashMap<>();
    
    // 存储分类订阅关系 - 分类ID -> 用户ID集合
    private static final Map<String, CopyOnWriteArraySet<String>> CATEGORY_SUBSCRIBERS = new ConcurrentHashMap<>();
    
    // 存储标签订阅关系 - 标签 -> 用户ID集合
    private static final Map<String, CopyOnWriteArraySet<String>> TAG_SUBSCRIBERS = new ConcurrentHashMap<>();
    
    // 存储会话信息 - 会话ID -> 用户信息
    private static final Map<String, UserSessionInfo> SESSION_INFO = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String userId = (String) session.getAttributes().get("userId");
        String categories = (String) session.getAttributes().get("categories");
        String tags = (String) session.getAttributes().get("tags");
        String clientIp = (String) session.getAttributes().get("clientIp");
        
        // 存储用户会话
        USER_SESSIONS.computeIfAbsent(userId, k -> new CopyOnWriteArraySet<>()).add(session);
        
        // 存储会话信息
        UserSessionInfo sessionInfo = new UserSessionInfo();
        sessionInfo.setUserId(userId);
        sessionInfo.setClientIp(clientIp);
        sessionInfo.setConnectTime(LocalDateTime.now());
        SESSION_INFO.put(session.getId(), sessionInfo);
        
        // 订阅分类
        if (StringUtils.hasText(categories)) {
            List<String> categoryList = Arrays.asList(categories.split(","));
            for (String categoryId : categoryList) {
                CATEGORY_SUBSCRIBERS.computeIfAbsent(categoryId, k -> new CopyOnWriteArraySet<>()).add(userId);
            }
        }
        
        // 订阅标签
        if (StringUtils.hasText(tags)) {
            List<String> tagList = Arrays.asList(tags.split(","));
            for (String tag : tagList) {
                TAG_SUBSCRIBERS.computeIfAbsent(tag, k -> new CopyOnWriteArraySet<>()).add(userId);
            }
        }
        
        log.info("新闻推送WebSocket连接建立：用户ID={}, 会话ID={}, IP={}", 
                userId, session.getId(), clientIp);
        
        // 发送连接成功消息
        sendMessage(session, createSystemMessage("CONNECTION_SUCCESS", "连接成功"));
        
        // 发送欢迎消息
        sendMessage(session, createSystemMessage("WELCOME", "欢迎使用实时新闻推送服务"));
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        if (!(message instanceof TextMessage)) {
            return;
        }
        
        String payload = ((TextMessage) message).getPayload();
        log.debug("收到WebSocket消息：{}", payload);
        
        try {
            WebSocketRequest wsRequest = objectMapper.readValue(payload, WebSocketRequest.class);
            UserSessionInfo sessionInfo = SESSION_INFO.get(session.getId());
            
            if (sessionInfo == null) {
                sendMessage(session, createErrorMessage("SESSION_NOT_FOUND", "会话信息不存在"));
                return;
            }
            
            String messageType = wsRequest.getType();
            
            switch (messageType) {
                case "SUBSCRIBE_CATEGORY":
                    handleSubscribeCategory(session, wsRequest, sessionInfo);
                    break;
                case "UNSUBSCRIBE_CATEGORY":
                    handleUnsubscribeCategory(session, wsRequest, sessionInfo);
                    break;
                case "SUBSCRIBE_TAG":
                    handleSubscribeTag(session, wsRequest, sessionInfo);
                    break;
                case "UNSUBSCRIBE_TAG":
                    handleUnsubscribeTag(session, wsRequest, sessionInfo);
                    break;
                case "MARK_READ":
                    handleMarkRead(session, wsRequest, sessionInfo);
                    break;
                case "HEARTBEAT":
                    handleHeartbeat(session, sessionInfo);
                    break;
                default:
                    sendMessage(session, createErrorMessage("UNKNOWN_MESSAGE_TYPE", "未知消息类型"));
            }
        } catch (Exception e) {
            log.error("处理WebSocket消息异常", e);
            sendMessage(session, createErrorMessage("MESSAGE_PARSE_ERROR", "消息解析失败"));
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输异常：会话ID={}", session.getId(), exception);
        cleanupSession(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        log.info("WebSocket连接关闭：会话ID={}, 状态={}", session.getId(), closeStatus);
        cleanupSession(session);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 处理订阅分类
     */
    private void handleSubscribeCategory(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        try {
            String categoryId = (String) wsRequest.getData();
            CATEGORY_SUBSCRIBERS.computeIfAbsent(categoryId, k -> new CopyOnWriteArraySet<>()).add(sessionInfo.getUserId());
            
            sendMessage(session, createSuccessMessage("CATEGORY_SUBSCRIBED", 
                       Map.of("categoryId", categoryId, "message", "分类订阅成功")));
            
            log.info("用户 {} 订阅了分类 {}", sessionInfo.getUserId(), categoryId);
        } catch (Exception e) {
            log.error("处理订阅分类异常", e);
            sendMessage(session, createErrorMessage("SUBSCRIBE_CATEGORY_ERROR", "订阅分类失败"));
        }
    }

    /**
     * 处理取消订阅分类
     */
    private void handleUnsubscribeCategory(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        try {
            String categoryId = (String) wsRequest.getData();
            CopyOnWriteArraySet<String> subscribers = CATEGORY_SUBSCRIBERS.get(categoryId);
            if (subscribers != null) {
                subscribers.remove(sessionInfo.getUserId());
            }
            
            sendMessage(session, createSuccessMessage("CATEGORY_UNSUBSCRIBED", 
                       Map.of("categoryId", categoryId, "message", "取消分类订阅成功")));
            
            log.info("用户 {} 取消订阅了分类 {}", sessionInfo.getUserId(), categoryId);
        } catch (Exception e) {
            log.error("处理取消订阅分类异常", e);
            sendMessage(session, createErrorMessage("UNSUBSCRIBE_CATEGORY_ERROR", "取消订阅分类失败"));
        }
    }

    /**
     * 处理订阅标签
     */
    private void handleSubscribeTag(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        try {
            String tag = (String) wsRequest.getData();
            TAG_SUBSCRIBERS.computeIfAbsent(tag, k -> new CopyOnWriteArraySet<>()).add(sessionInfo.getUserId());
            
            sendMessage(session, createSuccessMessage("TAG_SUBSCRIBED", 
                       Map.of("tag", tag, "message", "标签订阅成功")));
            
            log.info("用户 {} 订阅了标签 {}", sessionInfo.getUserId(), tag);
        } catch (Exception e) {
            log.error("处理订阅标签异常", e);
            sendMessage(session, createErrorMessage("SUBSCRIBE_TAG_ERROR", "订阅标签失败"));
        }
    }

    /**
     * 处理取消订阅标签
     */
    private void handleUnsubscribeTag(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        try {
            String tag = (String) wsRequest.getData();
            CopyOnWriteArraySet<String> subscribers = TAG_SUBSCRIBERS.get(tag);
            if (subscribers != null) {
                subscribers.remove(sessionInfo.getUserId());
            }
            
            sendMessage(session, createSuccessMessage("TAG_UNSUBSCRIBED", 
                       Map.of("tag", tag, "message", "取消标签订阅成功")));
            
            log.info("用户 {} 取消订阅了标签 {}", sessionInfo.getUserId(), tag);
        } catch (Exception e) {
            log.error("处理取消订阅标签异常", e);
            sendMessage(session, createErrorMessage("UNSUBSCRIBE_TAG_ERROR", "取消订阅标签失败"));
        }
    }

    /**
     * 处理标记已读
     */
    private void handleMarkRead(WebSocketSession session, WebSocketRequest wsRequest, UserSessionInfo sessionInfo) {
        try {
            String newsId = (String) wsRequest.getData();
            
            // 这里可以调用服务层方法记录用户阅读行为
            
            sendMessage(session, createSuccessMessage("NEWS_MARKED_READ", 
                       Map.of("newsId", newsId, "message", "标记已读成功")));
            
            log.debug("用户 {} 标记新闻 {} 为已读", sessionInfo.getUserId(), newsId);
        } catch (Exception e) {
            log.error("处理标记已读异常", e);
            sendMessage(session, createErrorMessage("MARK_READ_ERROR", "标记已读失败"));
        }
    }

    /**
     * 处理心跳
     */
    private void handleHeartbeat(WebSocketSession session, UserSessionInfo sessionInfo) {
        sessionInfo.setLastHeartbeatTime(LocalDateTime.now());
        sendMessage(session, createSystemMessage("HEARTBEAT_RESPONSE", "pong"));
    }

    /**
     * 向指定用户推送新闻
     */
    public void pushNewsToUser(String userId, Object newsData) {
        CopyOnWriteArraySet<WebSocketSession> sessions = USER_SESSIONS.get(userId);
        if (sessions != null) {
            sessions.forEach(session -> {
                if (session.isOpen()) {
                    sendMessage(session, createDataMessage("NEW_NEWS", newsData));
                }
            });
        }
    }

    /**
     * 向分类订阅者推送新闻
     */
    public void pushNewsToCategory(String categoryId, Object newsData) {
        CopyOnWriteArraySet<String> subscribers = CATEGORY_SUBSCRIBERS.get(categoryId);
        if (subscribers != null) {
            subscribers.forEach(userId -> pushNewsToUser(userId, newsData));
        }
    }

    /**
     * 向标签订阅者推送新闻
     */
    public void pushNewsToTag(String tag, Object newsData) {
        CopyOnWriteArraySet<String> subscribers = TAG_SUBSCRIBERS.get(tag);
        if (subscribers != null) {
            subscribers.forEach(userId -> pushNewsToUser(userId, newsData));
        }
    }

    /**
     * 广播新闻给所有在线用户
     */
    public void broadcastNews(Object newsData) {
        USER_SESSIONS.values().forEach(sessions -> 
            sessions.forEach(session -> {
                if (session.isOpen()) {
                    sendMessage(session, createDataMessage("BREAKING_NEWS", newsData));
                }
            })
        );
    }

    /**
     * 发送消息到WebSocket会话
     */
    private void sendMessage(WebSocketSession session, Object message) {
        if (session.isOpen()) {
            try {
                String json = objectMapper.writeValueAsString(message);
                session.sendMessage(new TextMessage(json));
            } catch (IOException e) {
                log.error("发送WebSocket消息异常", e);
            }
        }
    }

    /**
     * 清理会话
     */
    private void cleanupSession(WebSocketSession session) {
        UserSessionInfo sessionInfo = SESSION_INFO.remove(session.getId());
        if (sessionInfo != null) {
            String userId = sessionInfo.getUserId();
            
            // 移除用户会话
            CopyOnWriteArraySet<WebSocketSession> sessions = USER_SESSIONS.get(userId);
            if (sessions != null) {
                sessions.remove(session);
                if (sessions.isEmpty()) {
                    USER_SESSIONS.remove(userId);
                }
            }
            
            // 移除分类订阅
            CATEGORY_SUBSCRIBERS.values().forEach(subscribers -> subscribers.remove(userId));
            
            // 移除标签订阅
            TAG_SUBSCRIBERS.values().forEach(subscribers -> subscribers.remove(userId));
            
            log.info("清理WebSocket会话：用户ID={}, 会话ID={}", userId, session.getId());
        }
    }

    /**
     * 创建系统消息
     */
    private WebSocketResponse createSystemMessage(String type, String message) {
        WebSocketResponse response = new WebSocketResponse();
        response.setType(type);
        response.setMessage(message);
        response.setTimestamp(LocalDateTime.now());
        return response;
    }

    /**
     * 创建成功消息
     */
    private WebSocketResponse createSuccessMessage(String type, Object data) {
        WebSocketResponse response = new WebSocketResponse();
        response.setType(type);
        response.setSuccess(true);
        response.setData(data);
        response.setTimestamp(LocalDateTime.now());
        return response;
    }

    /**
     * 创建错误消息
     */
    private WebSocketResponse createErrorMessage(String type, String message) {
        WebSocketResponse response = new WebSocketResponse();
        response.setType(type);
        response.setSuccess(false);
        response.setMessage(message);
        response.setTimestamp(LocalDateTime.now());
        return response;
    }

    /**
     * 创建数据消息
     */
    private WebSocketResponse createDataMessage(String type, Object data) {
        WebSocketResponse response = new WebSocketResponse();
        response.setType(type);
        response.setSuccess(true);
        response.setData(data);
        response.setTimestamp(LocalDateTime.now());
        return response;
    }

    /**
     * 用户会话信息
     */
    private static class UserSessionInfo {
        private String userId;
        private String clientIp;
        private LocalDateTime connectTime;
        private LocalDateTime lastHeartbeatTime;

        // Getters and Setters
        public String getUserId() { return userId; }
        public void setUserId(String userId) { this.userId = userId; }
        
        public String getClientIp() { return clientIp; }
        public void setClientIp(String clientIp) { this.clientIp = clientIp; }
        
        public LocalDateTime getConnectTime() { return connectTime; }
        public void setConnectTime(LocalDateTime connectTime) { this.connectTime = connectTime; }
        
        public LocalDateTime getLastHeartbeatTime() { return lastHeartbeatTime; }
        public void setLastHeartbeatTime(LocalDateTime lastHeartbeatTime) { this.lastHeartbeatTime = lastHeartbeatTime; }
    }

    /**
     * WebSocket请求
     */
    private static class WebSocketRequest {
        private String type;
        private Object data;

        // Getters and Setters
        public String getType() { return type; }
        public void setType(String type) { this.type = type; }
        
        public Object getData() { return data; }
        public void setData(Object data) { this.data = data; }
    }

    /**
     * WebSocket响应
     */
    private static class WebSocketResponse {
        private String type;
        private Boolean success;
        private String message;
        private Object data;
        private LocalDateTime timestamp;

        // Getters and Setters
        public String getType() { return type; }
        public void setType(String type) { this.type = type; }
        
        public Boolean getSuccess() { return success; }
        public void setSuccess(Boolean success) { this.success = success; }
        
        public String getMessage() { return message; }
        public void setMessage(String message) { this.message = message; }
        
        public Object getData() { return data; }
        public void setData(Object data) { this.data = data; }
        
        public LocalDateTime getTimestamp() { return timestamp; }
        public void setTimestamp(LocalDateTime timestamp) { this.timestamp = timestamp; }
    }
}
